# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: donggu

import sys
import neptune
import re
import optuna
import neptunecontrib.monitoring.optuna as opt_utils

name_filter = sys.argv[1]

session = neptune.sessions.Session.with_default_backend().get_project('donggu-prestolabs/lm-agg2')
exps = session.get_experiments(state=['failed', 'running'], tag=[name_filter])

for exp in exps:
  if exp.name != name_filter:
    print('skipping ', exp.id, exp.name)
    continue
  print('processing ', exp.id, exp.name)

  prop = exp.get_properties()
  exchange = prop['exchange']
  target = prop['target']
  name = f'{exp.name}-{exchange}-{target}'

  # load local
  study = optuna.load_study(study_name=name,
                            storage=f'postgresql://postgres:EU2kb6Xos9pS@10.21.105.13:5432')
  opt_utils.log_study_info(study, experiment=exp)

